{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "VBefPIvx73o7",
        "outputId": "d2336aed-1a1d-4af7-ac85-f71c514c3885"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Cloning into 'PyEPO'...\n",
            "remote: Enumerating objects: 126, done.\u001b[K\n",
            "remote: Counting objects: 100% (126/126), done.\u001b[K\n",
            "remote: Compressing objects: 100% (111/111), done.\u001b[K\n",
            "remote: Total 126 (delta 22), reused 65 (delta 12), pack-reused 0\u001b[K\n",
            "Receiving objects: 100% (126/126), 2.44 MiB | 19.67 MiB/s, done.\n",
            "Resolving deltas: 100% (22/22), done.\n",
            "Processing ./PyEPO/pkg\n",
            "  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from pyepo==0.3.4) (1.23.5)\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.10/dist-packages (from pyepo==0.3.4) (1.10.1)\n",
            "Collecting pathos (from pyepo==0.3.4)\n",
            "  Downloading pathos-0.3.1-py3-none-any.whl (82 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m82.1/82.1 kB\u001b[0m \u001b[31m436.7 kB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: tqdm in /usr/local/lib/python3.10/dist-packages (from pyepo==0.3.4) (4.66.1)\n",
            "Collecting Pyomo>=6.1.2 (from pyepo==0.3.4)\n",
            "  Downloading Pyomo-6.6.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (11.9 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m11.9/11.9 MB\u001b[0m \u001b[31m58.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting gurobipy>=9.1.2 (from pyepo==0.3.4)\n",
            "  Downloading gurobipy-10.0.2-cp310-cp310-manylinux2014_x86_64.whl (12.7 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m12.7/12.7 MB\u001b[0m \u001b[31m57.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: scikit_learn in /usr/local/lib/python3.10/dist-packages (from pyepo==0.3.4) (1.2.2)\n",
            "Requirement already satisfied: torch>=1.13.1 in /usr/local/lib/python3.10/dist-packages (from pyepo==0.3.4) (2.0.1+cu118)\n",
            "Collecting ply (from Pyomo>=6.1.2->pyepo==0.3.4)\n",
            "  Downloading ply-3.11-py2.py3-none-any.whl (49 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m49.6/49.6 kB\u001b[0m \u001b[31m5.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from torch>=1.13.1->pyepo==0.3.4) (3.12.2)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.10/dist-packages (from torch>=1.13.1->pyepo==0.3.4) (4.7.1)\n",
            "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch>=1.13.1->pyepo==0.3.4) (1.12)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch>=1.13.1->pyepo==0.3.4) (3.1)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch>=1.13.1->pyepo==0.3.4) (3.1.2)\n",
            "Requirement already satisfied: triton==2.0.0 in /usr/local/lib/python3.10/dist-packages (from torch>=1.13.1->pyepo==0.3.4) (2.0.0)\n",
            "Requirement already satisfied: cmake in /usr/local/lib/python3.10/dist-packages (from triton==2.0.0->torch>=1.13.1->pyepo==0.3.4) (3.27.2)\n",
            "Requirement already satisfied: lit in /usr/local/lib/python3.10/dist-packages (from triton==2.0.0->torch>=1.13.1->pyepo==0.3.4) (16.0.6)\n",
            "Collecting ppft>=1.7.6.7 (from pathos->pyepo==0.3.4)\n",
            "  Downloading ppft-1.7.6.7-py3-none-any.whl (56 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m56.8/56.8 kB\u001b[0m \u001b[31m4.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting dill>=0.3.7 (from pathos->pyepo==0.3.4)\n",
            "  Downloading dill-0.3.7-py3-none-any.whl (115 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m115.3/115.3 kB\u001b[0m \u001b[31m13.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting pox>=0.3.3 (from pathos->pyepo==0.3.4)\n",
            "  Downloading pox-0.3.3-py3-none-any.whl (29 kB)\n",
            "Collecting multiprocess>=0.70.15 (from pathos->pyepo==0.3.4)\n",
            "  Downloading multiprocess-0.70.15-py310-none-any.whl (134 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m134.8/134.8 kB\u001b[0m \u001b[31m13.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: joblib>=1.1.1 in /usr/local/lib/python3.10/dist-packages (from scikit_learn->pyepo==0.3.4) (1.3.2)\n",
            "Requirement already satisfied: threadpoolctl>=2.0.0 in /usr/local/lib/python3.10/dist-packages (from scikit_learn->pyepo==0.3.4) (3.2.0)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch>=1.13.1->pyepo==0.3.4) (2.1.3)\n",
            "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.10/dist-packages (from sympy->torch>=1.13.1->pyepo==0.3.4) (1.3.0)\n",
            "Building wheels for collected packages: pyepo\n",
            "  Building wheel for pyepo (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for pyepo: filename=pyepo-0.3.4-py3-none-any.whl size=40269 sha256=b97496311d68718030a1a9a8697e78afbe09a9e3321044961919ae373e68bb41\n",
            "  Stored in directory: /tmp/pip-ephem-wheel-cache-t9vmxfgw/wheels/46/e0/92/19132b049913f800d1a8d6e61b81fe00bcac3f9d11ef30d2a5\n",
            "Successfully built pyepo\n",
            "Installing collected packages: ply, gurobipy, Pyomo, ppft, pox, dill, multiprocess, pathos, pyepo\n",
            "Successfully installed Pyomo-6.6.1 dill-0.3.7 gurobipy-10.0.2 multiprocess-0.70.15 pathos-0.3.1 ply-3.11 pox-0.3.3 ppft-1.7.6.7 pyepo-0.3.4\n",
            "Collecting coptpy\n",
            "  Downloading coptpy-6.5.8-cp310-cp310-manylinux2014_x86_64.whl (10.5 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m10.5/10.5 MB\u001b[0m \u001b[31m66.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hInstalling collected packages: coptpy\n",
            "Successfully installed coptpy-6.5.8\n"
          ]
        }
      ],
      "source": [
        "# download\n",
        "!git clone -b main --depth 1 https://github.com/khalil-research/PyEPO.git\n",
        "# install\n",
        "!pip install PyEPO/pkg/.\n",
        "!pip install coptpy"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HjxG2TL58ZFb"
      },
      "source": [
        "## Build optModel"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "nRbnYdWx8M8L",
        "outputId": "ec2ca202-9f1a-43d1-dc2b-0977d7d1abec"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Auto-Sklearn cannot be imported.\n",
            "Cardinal Optimizer v6.5.8. Build date Aug 23 2023\n",
            "Copyright Cardinal Operations 2023. All Rights Reserved\n",
            "\n"
          ]
        }
      ],
      "source": [
        "from coptpy import COPT\n",
        "from coptpy import Envr\n",
        "from pyepo.model.copt import optCoptModel\n",
        "\n",
        "class myOptModel(optCoptModel):\n",
        "    def _getModel(self):\n",
        "        # ceate a model\n",
        "        m = Envr().createModel()\n",
        "        # varibles\n",
        "        x = m.addVars(5, nameprefix='x', vtype=COPT.BINARY)\n",
        "        # sense\n",
        "        m.setObjSense(COPT.MAXIMIZE)\n",
        "        # constraints\n",
        "        m.addConstr(3*x[0]+4*x[1]+3*x[2]+6*x[3]+4*x[4]<=12)\n",
        "        m.addConstr(4*x[0]+5*x[1]+2*x[2]+3*x[3]+5*x[4]<=10)\n",
        "        m.addConstr(5*x[0]+4*x[1]+6*x[2]+2*x[3]+3*x[4]<=15)\n",
        "        return m, x\n",
        "\n",
        "optmodel = myOptModel()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-vsQ0Vre84_1"
      },
      "source": [
        "## Problem Data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "-rqtFovN8-Gc"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "torch.manual_seed(42)\n",
        "\n",
        "num_data = 1000 # number of data\n",
        "num_feat = 5 # feature dimention\n",
        "num_cost = 5 # cost dimention\n",
        "\n",
        "# randomly generate data\n",
        "x_true = torch.rand(num_data, num_feat) # feature\n",
        "weight_true = torch.rand(num_feat, num_cost) # weight\n",
        "bias_true = torch.randn(num_cost) # bias\n",
        "noise = 0.5 * torch.randn(num_data, num_cost) # random noise\n",
        "c_true = x_true @ weight_true + bias_true + noise # cost coef"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Tk_wcRiq_5_K",
        "outputId": "1b73834c-421b-4bbe-e5c3-385d7e963445"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Optimizing for optDataset...\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 800/800 [00:05<00:00, 147.66it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Optimizing for optDataset...\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 200/200 [00:01<00:00, 108.59it/s]\n"
          ]
        }
      ],
      "source": [
        "# split train test data\n",
        "from sklearn.model_selection import train_test_split\n",
        "x_train, x_test, c_train, c_test = train_test_split(x_true, c_true, test_size=200, random_state=42)\n",
        "\n",
        "# build optDataset\n",
        "from pyepo.data.dataset import optDataset\n",
        "dataset_train = optDataset(optmodel, x_train, c_train)\n",
        "dataset_test = optDataset(optmodel, x_test, c_test)\n",
        "\n",
        "# build DataLoader\n",
        "from torch.utils.data import DataLoader\n",
        "batch_size = 32\n",
        "loader_train = DataLoader(dataset_train, batch_size=batch_size, shuffle=True)\n",
        "loader_test = DataLoader(dataset_test, batch_size=batch_size, shuffle=False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZENnTDsKEwho"
      },
      "source": [
        "## Build Prediction Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "g2mNrexpEvhS"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "from torch import nn\n",
        "\n",
        "# build linear model\n",
        "class LinearRegression(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(LinearRegression, self).__init__()\n",
        "        self.linear = nn.Linear(num_feat, num_cost)\n",
        "\n",
        "    def forward(self, x):\n",
        "        out = self.linear(x)\n",
        "        return out\n",
        "\n",
        "# init model\n",
        "reg = LinearRegression()\n",
        "# cuda\n",
        "if torch.cuda.is_available():\n",
        "    reg = reg.cuda()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oyDZry1AD99r"
      },
      "source": [
        "## AutoGrad Module for Optimization"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4Ykku_3eDsNc",
        "outputId": "c92e76d2-099f-48ca-d43c-e6ef0a4adcbb"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Num of cores: 2\n",
            "Num of cores: 2\n",
            "Num of cores: 2\n"
          ]
        }
      ],
      "source": [
        "import pyepo\n",
        "\n",
        "# init SPO+ loss\n",
        "spop = pyepo.func.SPOPlus(optmodel, processes=2)\n",
        "# init PFY loss\n",
        "pfy = pyepo.func.perturbedFenchelYoung(optmodel, n_samples=3, sigma=1.0, processes=2)\n",
        "# init NCE loss\n",
        "nce = pyepo.func.NCE(optmodel, processes=2, solve_ratio=0.05, dataset=dataset_train)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Ets2ndfSFnag",
        "outputId": "c1f00803-a7f1-42cd-f437-46c6e6faf1cb"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Loss:    6.5576,  Regret: 21.4272%\n",
            "Loss:    3.9369,  Regret:  0.1530%\n",
            "Loss:    2.1917,  Regret:  0.1530%\n",
            "Loss:    0.7715,  Regret:  0.1530%\n",
            "Loss:    0.4534,  Regret:  0.1530%\n"
          ]
        }
      ],
      "source": [
        "  # set adam optimizer\n",
        "  optimizer = torch.optim.Adam(reg.parameters(), lr=5e-3)\n",
        "\n",
        "  # train mode\n",
        "  reg.train()\n",
        "  for epoch in range(5):\n",
        "    # load data\n",
        "    for i, data in enumerate(loader_train):\n",
        "        x, c, w, z = data # feat, cost, sol, obj\n",
        "        # cuda\n",
        "        if torch.cuda.is_available():\n",
        "            x, c, w, z = x.cuda(), c.cuda(), w.cuda(), z.cuda()\n",
        "        # forward pass\n",
        "        cp = reg(x)\n",
        "        loss = spop(cp, c, w, z)\n",
        "        # backward pass\n",
        "        optimizer.zero_grad()\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "    # log\n",
        "    regret = pyepo.metric.regret(reg, optmodel, loader_test)\n",
        "    print(\"Loss: {:9.4f},  Regret: {:7.4f}%\".format(loss.item(), regret*100))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "yBl4rTYk1vEf"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.9.15"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}